﻿import pygame
import sys
import random
import math
from pygame.locals import *

# 初始化 Pygame
pygame.init()

# 屏幕设置
WIDTH, HEIGHT = 1000, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("堆垛机搬运原理展示")

# 颜色定义
BACKGROUND = (30, 35, 45)
SHELF_COLOR = (100, 110, 130)
SHELF_HIGHLIGHT = (150, 160, 180)
CRANE_COLOR = (200, 50, 50)
CRANE_ACCENT = (250, 150, 50)
FORK_COLOR = (70, 170, 220)
PATH_COLOR = (100, 200, 100)
TEXT_COLOR = (220, 220, 220)
HIGHLIGHT = (255, 215, 0)
GRID_COLOR = (60, 70, 90)
ITEM_COLOR = (70, 170, 220)
ITEM_HIGHLIGHT = (100, 200, 250)

# 字体
title_font = pygame.font.SysFont("simhei", 36, bold=True)
text_font = pygame.font.SysFont("simhei", 24)
small_font = pygame.font.SysFont("simhei", 18)

# 时钟
clock = pygame.time.Clock()
FPS = 60

# 仓库配置
NUM_SHELVES = 10
SHELF_WIDTH = 80
SHELF_HEIGHT = 150
SHELF_DEPTH = 60
AISLE_WIDTH = 120
FLOOR_HEIGHT = 50

# 堆垛机配置
CRANE_WIDTH = 40
CRANE_HEIGHT = 100
FORK_WIDTH = 20
FORK_HEIGHT = SHELF_DEPTH
FORK_SPEED = 3
CRANE_SPEED = 4

class Warehouse:
    def __init__(self):
        self.shelves = []
        self.items = []
        self.crane = StackerCrane()
        self.target_item = None
        self.target_shelf = None
        self.mode = "idle"  # idle, moving_to_item, picking, moving_to_shelf, placing
        self.create_shelves()
        self.create_items()
        self.path = []
        self.instructions = [
            "堆垛机搬运原理展示",
            "按空格键开始/暂停模拟",
            "按R键重置模拟",
            "按1-3键选择不同搬运任务"
        ]
    
    def create_shelves(self):
        # 创建货架
        for i in range(NUM_SHELVES):
            # 左右两侧货架
            left_x = 100
            right_x = WIDTH - 100 - SHELF_WIDTH
            y = 150 + i * (SHELF_HEIGHT + 20)
            
            # 左侧货架
            self.shelves.append({
                "rect": pygame.Rect(left_x, y, SHELF_WIDTH, SHELF_HEIGHT),
                "side": "left",
                "items": []
            })
            
            # 右侧货架
            self.shelves.append({
                "rect": pygame.Rect(right_x, y, SHELF_WIDTH, SHELF_HEIGHT),
                "side": "right",
                "items": []
            })
    
    def create_items(self):
        # 在货架上创建货物
        for shelf in self.shelves:
            for _ in range(3):  # 每个货架3个货物
                item_x = shelf["rect"].x + random.randint(10, SHELF_WIDTH - 30)
                item_y = shelf["rect"].y + random.randint(10, SHELF_HEIGHT - 30)
                self.items.append({
                    "rect": pygame.Rect(item_x, item_y, 20, 20),
                    "shelf": shelf,
                    "color": ITEM_COLOR
                })
                shelf["items"].append(self.items[-1])
    
    def find_item(self):
        # 随机选择一个货物作为目标
        if not self.items:
            return None
        return random.choice(self.items)
    
    def find_empty_shelf(self):
        # 随机找一个空货位
        for shelf in self.shelves:
            if len(shelf["items"]) < 3:  # 每个货架最多3个货物
                return shelf
        return None
    
    def start_task(self, task_type=1):
        self.mode = "moving_to_item"
        
        if task_type == 1:
            # 任务1: 从货架取货放到另一个货架
            self.target_item = self.find_item()
            if self.target_item:
                self.target_shelf = self.find_empty_shelf()
                if not self.target_shelf:
                    self.target_shelf = random.choice(self.shelves)
        elif task_type == 2:
            # 任务2: 将货物从货架移动到地面
            self.target_item = self.find_item()
            self.target_shelf = None
        elif task_type == 3:
            # 任务3: 将货物从地面移动到货架
            self.target_item = None
            self.target_shelf = self.find_empty_shelf()
            if not self.target_shelf:
                self.target_shelf = random.choice(self.shelves)
            
            # 创建一个地面货物
            ground_x = WIDTH // 2 + random.randint(-100, 100)
            ground_y = HEIGHT - FLOOR_HEIGHT - 30
            self.items.append({
                "rect": pygame.Rect(ground_x, ground_y, 20, 20),
                "shelf": None,
                "color": ITEM_HIGHLIGHT
            })
            self.target_item = self.items[-1]
    
    def update(self):
        if self.mode == "idle":
            return
        
        # 更新堆垛机状态
        crane = self.crane
        
        if self.mode == "moving_to_item":
            # 移动到货物位置
            if self.target_item:
                target_x = self.target_item["rect"].centerx
                target_y = self.target_item["rect"].centery - 20
                
                # 计算移动方向
                dx = target_x - crane.x
                dy = target_y - crane.y
                distance = math.sqrt(dx*dx + dy*dy)
                
                if distance < CRANE_SPEED:
                    crane.x = target_x
                    crane.y = target_y
                    self.mode = "picking"
                else:
                    crane.x += dx / distance * CRANE_SPEED
                    crane.y += dy / distance * CRANE_SPEED
            else:
                self.mode = "idle"
        
        elif self.mode == "picking":
            # 伸出货叉取货
            if crane.fork_position < FORK_HEIGHT:
                crane.fork_position += FORK_SPEED
            else:
                # 取货完成
                if self.target_item in self.items:
                    self.items.remove(self.target_item)
                    if self.target_item["shelf"]:
                        self.target_item["shelf"]["items"].remove(self.target_item)
                    crane.carrying = self.target_item
                
                if self.target_shelf:
                    self.mode = "moving_to_shelf"
                else:
                    self.mode = "placing"
        
        elif self.mode == "moving_to_shelf":
            # 移动到目标货架位置
            if self.target_shelf:
                target_x = self.target_shelf["rect"].centerx
                target_y = self.target_shelf["rect"].centery - 20
                
                # 计算移动方向
                dx = target_x - crane.x
                dy = target_y - crane.y
                distance = math.sqrt(dx*dx + dy*dy)
                
                if distance < CRANE_SPEED:
                    crane.x = target_x
                    crane.y = target_y
                    self.mode = "placing"
                else:
                    crane.x += dx / distance * CRANE_SPEED
                    crane.y += dy / distance * CRANE_SPEED
        
        elif self.mode == "placing":
            # 放置货物
            if crane.fork_position > 0:
                crane.fork_position -= FORK_SPEED
            else:
                # 放置货物完成
                if crane.carrying:
                    crane.carrying["rect"].center = (crane.x, crane.y + crane.height + 10)
                    if self.target_shelf:
                        crane.carrying["shelf"] = self.target_shelf
                        self.target_shelf["items"].append(crane.carrying)
                    self.items.append(crane.carrying)
                    crane.carrying = None
                
                self.mode = "idle"
                self.target_item = None
                self.target_shelf = None
    
    def draw(self, surface):
        # 绘制背景
        surface.fill(BACKGROUND)
        
        # 绘制标题
        title = title_font.render("堆垛机搬运原理展示", True, TEXT_COLOR)
        surface.blit(title, (WIDTH // 2 - title.get_width() // 2, 20))
        
        # 绘制地面
        pygame.draw.rect(surface, (50, 55, 65), (0, HEIGHT - FLOOR_HEIGHT, WIDTH, FLOOR_HEIGHT))
        pygame.draw.line(surface, (70, 75, 85), (0, HEIGHT - FLOOR_HEIGHT), (WIDTH, HEIGHT - FLOOR_HEIGHT), 2)
        
        # 绘制货架
        for shelf in self.shelves:
            # 绘制货架主体
            pygame.draw.rect(surface, SHELF_COLOR, shelf["rect"])
            pygame.draw.rect(surface, SHELF_HIGHLIGHT, shelf["rect"], 2)
            
            # 绘制货架层板
            for i in range(1, 4):
                y_pos = shelf["rect"].y + i * (SHELF_HEIGHT // 4)
                pygame.draw.line(surface, SHELF_HIGHLIGHT, 
                                (shelf["rect"].x, y_pos), 
                                (shelf["rect"].x + shelf["rect"].width, y_pos), 2)
        
        # 绘制货物
        for item in self.items:
            pygame.draw.rect(surface, item["color"], item["rect"])
            pygame.draw.rect(surface, ITEM_HIGHLIGHT, item["rect"], 1)
        
        # 绘制堆垛机
        self.crane.draw(surface)
        
        # 绘制操作说明
        for i, instruction in enumerate(self.instructions):
            text = text_font.render(instruction, True, TEXT_COLOR)
            surface.blit(text, (20, 70 + i * 30))
        
        # 绘制当前模式
        mode_text = text_font.render(f"当前模式: {self.get_mode_text()}", True, HIGHLIGHT)
        surface.blit(mode_text, (WIDTH - mode_text.get_width() - 20, 70))
        
        # 绘制任务说明
        if self.target_item:
            item_text = small_font.render(f"目标货物: #{self.items.index(self.target_item)}", True, HIGHLIGHT)
            surface.blit(item_text, (WIDTH - item_text.get_width() - 20, 100))
        
        if self.target_shelf:
            shelf_idx = self.shelves.index(self.target_shelf)
            shelf_text = small_font.render(f"目标货架: #{shelf_idx}", True, HIGHLIGHT)
            surface.blit(shelf_text, (WIDTH - shelf_text.get_width() - 20, 130))
        
        # 绘制统计数据
        stats = [
            f"货物总数: {len(self.items)}",
            f"堆垛机状态: {'工作中' if self.mode != 'idle' else '空闲'}",
            f"搬运次数: {self.crane.task_count}"
        ]
        
        for i, stat in enumerate(stats):
            text = small_font.render(stat, True, TEXT_COLOR)
            surface.blit(text, (20, HEIGHT - 100 + i * 25))
    
    def get_mode_text(self):
        modes = {
            "idle": "空闲",
            "moving_to_item": "前往货物位置",
            "picking": "取货中",
            "moving_to_shelf": "前往货架位置",
            "placing": "放置货物"
        }
        return modes.get(self.mode, "未知状态")

class StackerCrane:
    def __init__(self):
        self.x = WIDTH // 2
        self.y = HEIGHT - FLOOR_HEIGHT - 50
        self.width = CRANE_WIDTH
        self.height = CRANE_HEIGHT
        self.fork_position = 0  # 货叉伸缩位置
        self.carrying = None  # 当前携带的货物
        self.task_count = 0
    
    def draw(self, surface):
        # 绘制堆垛机主体
        crane_rect = pygame.Rect(self.x - self.width // 2, 
                                self.y - self.height, 
                                self.width, self.height)
        pygame.draw.rect(surface, CRANE_COLOR, crane_rect)
        pygame.draw.rect(surface, CRANE_ACCENT, crane_rect, 2)
        
        # 绘制堆垛机顶部
        top_rect = pygame.Rect(self.x - self.width // 2 - 10, 
                              self.y - self.height - 10, 
                              self.width + 20, 10)
        pygame.draw.rect(surface, CRANE_ACCENT, top_rect)
        
        # 绘制货叉
        if self.fork_position > 0:
            fork_x = self.x
            fork_y = self.y - self.fork_position
            fork_rect = pygame.Rect(fork_x - FORK_WIDTH // 2, 
                                   fork_y, 
                                   FORK_WIDTH, FORK_HEIGHT)
            pygame.draw.rect(surface, FORK_COLOR, fork_rect)
            pygame.draw.rect(surface, CRANE_ACCENT, fork_rect, 2)
            
            # 绘制货叉与主体的连接线
            pygame.draw.line(surface, FORK_COLOR, 
                            (self.x, self.y), 
                            (fork_x, fork_y), 3)
        
        # 绘制携带的货物
        if self.carrying:
            cargo_rect = pygame.Rect(self.x - 15, self.y - self.fork_position - 20, 30, 30)
            pygame.draw.rect(surface, ITEM_HIGHLIGHT, cargo_rect)
            pygame.draw.rect(surface, HIGHLIGHT, cargo_rect, 2)

# 主函数
def main():
    warehouse = Warehouse()
    paused = False
    
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == QUIT:
                running = False
            elif event.type == KEYDOWN:
                if event.key == K_SPACE:
                    paused = not paused
                elif event.key == K_r:
                    warehouse = Warehouse()
                    paused = False
                elif event.key == K_1:
                    warehouse.start_task(1)
                elif event.key == K_2:
                    warehouse.start_task(2)
                elif event.key == K_3:
                    warehouse.start_task(3)
        
        # 更新仓库状态
        if not paused:
            warehouse.update()
        
        # 绘制
        warehouse.draw(screen)
        
        # 更新显示
        pygame.display.flip()
        
        # 控制帧率
        clock.tick(FPS)
    
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()